import init_file
import ROSinfo
import rospy
import math
import threading
from autopilot import navigate_to_target
from enu_xyz import GNC_COORDINATE
from queue import Queue
from collections import deque
import numpy as np
import torch
import time
import random
from TD3 import TD3
from TD3 import OUNoise

# 在全局范围内创建一个队列
boat_pos_queue = Queue()
boat_yaw_queue = Queue()
boat_vel_queue = Queue()
boat_speed_queue = Queue()
M_PI = math.pi/3
def GotoaPoint(given_point):
    while 1:
        boat_x, boat_y = boat_pos_queue.get()
        yaw = boat_yaw_queue.get()
        speed = boat_speed_queue.get()
        if ((boat_x - given_point[0]) ** 2 + (boat_y - given_point[1]) ** 2) < 100:
            print('end')
            break
        else:
            turn, power_command = navigate_to_target(yaw, boat_x, boat_y, given_point[0], given_point[1], speed)
            forward_v, lateral_v = boat_vel_queue.get()
            distance = math.sqrt((boat_x - given_point[0]) ** 2 + (boat_y - given_point[1]) ** 2)
            power_R = power_command
            power_L = power_R
            rosvrx_control.manual_act(power_L, power_R, turn)
        print('当前位置：', boat_x, boat_y, '当前朝向的点：', given_point, '当前距离：', distance)



def publish_commands():
    while not rospy.is_shutdown():
        state_dim = 6  # 根据实际情况设置状态维度
        action_dim = 1  # 根据实际情况设置动作维度
        hidden_dim = 128  # 根据实际情况设置隐藏层维度
        ddpg = TD3(action_dim, state_dim, hidden_dim)
        # ddpg.policy_net.load_state_dict(torch.load('/home/ltp/catkin_ws/src/deeplearning/scripts/net-weight.pt'))
        model = ddpg.policy_net
        USV_action_space = init_file.USV_ActionSpace()
        ou_noise = OUNoise(USV_action_space)
        max_episodes = 500
        max_steps = 500
        frame_idx = 0
        rewards = []
        turn_values = deque(maxlen=20)
        batch_size = 256
        episode = 0
        given_point = points[0].copy()
        new_point = points[1].copy()
        while episode < max_episodes:
            boat_x, boat_y = boat_pos_queue.get()
            yaw = boat_yaw_queue.get()
            forward_v, lateral_v = boat_vel_queue.get()
            random_num = random.uniform(-10, 10)
            print(random_num)
            if(boat_x > 650):
                given_point= points[1].copy()
                new_point = points[0].copy()
                GotoaPoint(given_point)
            elif boat_x < 0:
                given_point = points[0].copy()
                new_point = points[1].copy()
                GotoaPoint(given_point)
            else:
                t_error, ang_error = init_file.calculate_metrics(yaw, forward_v, lateral_v, points[0][0],
                                                                 points[0][1], points[1][0], points[1][1], boat_x, boat_y)
                if abs(t_error) > 15:
                    start_point = init_file.get_foot(points[0], points[1], boat_x, boat_y)
                    GotoaPoint(start_point)
                elif abs(t_error)<5:
                    given_point = [points[0][0], points[0][1] + random_num]
                    new_point = [points[1][0], points[1][1] + random_num]

            time.sleep(2)
            boat_x, boat_y = boat_pos_queue.get()
            yaw = boat_yaw_queue.get()
            forward_v, lateral_v = boat_vel_queue.get()

            t_error, ang_error = init_file.calculate_metrics(yaw, forward_v, lateral_v, given_point[0], given_point[1],
                                                           new_point[0], new_point[1], boat_x, boat_y)
            action = [0] * action_dim
            dt_t_error = 0
            pre_t_error = 0
            dt_ang_error = 0
            pre_ang_error = 0
            pre_time = time.time()
            state = np.array([t_error, ang_error, dt_t_error, dt_ang_error, forward_v] + action)
            ou_noise.reset()
            episode_reward = 0
            done = 0

            print(given_point,new_point)
            for step in range(max_steps):
                action = ddpg.policy_net.get_action(state)  # 根据当前的actor网络->策略选择动作
                noise_action , noise = ou_noise.get_action(action, frame_idx)  # 对动作添加噪音
                noise_action = noise_action.reshape(-1)
                turn = noise_action[0]
                power_L,power_R = init_file.servo_to_power(turn)
                turn_values.append(turn)
                rosvrx_control.manual_act(power_L, power_R, turn)
                boat_x, boat_y = boat_pos_queue.get()
                yaw = boat_yaw_queue.get()
                forward_v, lateral_v = boat_vel_queue.get()
                t_error, ang_error = init_file.calculate_metrics(yaw, forward_v, lateral_v, given_point[0], given_point[1],
                                                           new_point[0], new_point[1], boat_x, boat_y)
                current_time = time.time()
                dt_t_error = (t_error - pre_t_error) / (current_time - pre_time)
                dt_ang_error = (ang_error - pre_ang_error) / (current_time - pre_time)
                reward, done = init_file.reward_function(t_error, ang_error, turn_values, forward_v, done)
                next_state = np.concatenate(([t_error, ang_error, forward_v, dt_t_error, dt_ang_error], noise_action))
                ddpg.replay_buffer.push(state, noise_action, reward, next_state, done)  # 交互数据存储到经验缓存器
                # 从经验缓存器R中随机采样128个mini_batch数据(s_i,a_i,r_i,s_(i+1))
                if len(ddpg.replay_buffer) > batch_size:  # 根据经验缓存器的大小，随机采样数据并用于更新 DDPG 网络
                    ddpg.ddpg_update()  # 更新网络
                    # policy_loss_list.append(policy_loss), value_loss_list.append(value_loss)
                episode_reward += reward
                frame_idx += 1
                # if frame_idx % max(1000, max_steps + 1) == 0:
                #     plot(frame_idx, rewards)
                pre_t_error = t_error
                pre_ang_error = ang_error
                pre_time = current_time
                state = next_state
                if step % 20 == 0:
                    print('回合数:', episode, '步数:', step, '奖励:', reward, '动作:', action, '噪声后的动作:',
                      noise_action, '状态:', state)
                if done:
                    break
            episode += 1
            ddpg.cosine_annealing_lr(ddpg.policy_optimizer, episode, max_episodes)
            print('总奖励：', episode_reward)
            data = "Episode {}: {}".format(episode, episode_reward)
            init_file.write_data_to_file(data, 'data_origin.txt')
            rewards.append(episode_reward)
            torch.save(model.state_dict(), '/home/ltp/catkin_ws/src/deeplearning/scripts/net-weight.pt')
        init_file.plot(frame_idx, rewards)
        break


if __name__ == "__main__":

    gnc_coordinate = GNC_COORDINATE()

    # hangdian1_lon = 150.6749317
    # hangdian1_lat = -33.7222442
    # hangdian2_lon = 150.6759145
    # hangdian2_lat = -33.7226047
    # hangdian3_lon = 150.6763608
    # hangdian3_lat = -33.7219765
    # hangdian4_lon = 150.6754145
    # hangdian4_lat = -33.7216427

    # 训练路径航点
    points = [
        [150.6755737, -33.7221389],  # 第一个坐标点
        [150.6825302, -33.7224315],  # 第二个坐标点
    ]
    # 设置初始点经纬度
    origin_lon = points[0][0]
    origin_lat = points[0][1]
    #写一个for循环，将每个点的经纬度转换为坐标系
    for i in range(len(points)):
        points[i][0] = gnc_coordinate.Lon_trans_to_E(points[i][0], points[i][1], origin_lon, origin_lat)
        points[i][1] = gnc_coordinate.Lat_trans_to_N(points[i][0], points[i][1], origin_lon, origin_lat)


    print("points = [")
    for sublist in points:
        print("    [", end="")
        for item in sublist:
            print(item, end=", ")
        print("],")
    print("]")


    def update_func_gps(handle, lat, lon, alt):
        # 将ROS传回的经纬度转换为坐标系
        boat_x = gnc_coordinate.Lon_trans_to_E(lon, lat, origin_lon, origin_lat)
        boat_y = gnc_coordinate.Lat_trans_to_N(lon, lat, origin_lon, origin_lat)
        boat_pos_queue.put((boat_x, boat_y))


    def update_func_ea(handle, yaw, pitch, roll):
        boat_yaw_queue.put(yaw)

    def update_func_ground_speed(handle, speed , forward_v , lateral_v):
        boat_vel_queue.put((forward_v,lateral_v))
        boat_speed_queue.put(speed)

    def update_func_gyro(handle, gyroX, gyroY, gyroZ):
        pass

    def update_func_acc(handle, accX, accY, accZ):
        pass

    rospy.init_node('my_ROSVRX_node')

    update_handle = None

    ros_vrx_info = ROSinfo.ROSVRXInfo(update_func_gps, update_func_ea, update_func_ground_speed,
                              update_func_gyro, update_func_acc, update_handle)

    propeller_mid = 0
    servo_mid = 0
    rosvrx_control = ROSinfo.propulsion_VRX(propeller_mid, servo_mid)


    # 创建一个线程来发布命令
    publish_thread = threading.Thread(target=publish_commands)
    publish_thread.start()

    try:
        ros_vrx_info.ROSVRX_info_create()

    except rospy.ROSInterruptException:
        pass
